from openai import OpenAI


class DeepSeek():

    def __init__(self) -> None:
        self.client = OpenAI(
            api_key="ollama",
            base_url="http://localhost:11434/v1",
        )
        self.system_prompt = """
            用户将给你一段病例文本，你需要从中提取 '用户性别'，'症状'，'出生日期'。
            最后以JSON数组的形式输出，格式如下：
            {
                "list": [
                    {"key": "性别", "value": "<性别 没有请填 null>"},
                    {"key": "症状", "value": "<症状 没有请填 null>"},
                    {"key": "生日", "value": "<格式为 YYYY-mm-dd 没有请填 null>"}
                ]
            }
            如果用户没有提供相关信息，就用未查询到代替。
        """
        self.chat_prompt = "You are a helpful assistant."

    def getText(self, messages):
        response = self.client.chat.completions.create(
            model="deepseek-r1:1.5b",
            messages=[
                {'role': 'system', 'content': "this is a message"},
                {'role': 'user', 'content': messages}
            ]
        )
        return response.choices[0].message.content


if __name__ == "__main__":
    ds = DeepSeek()
    print(ds.getText("Hello"))
